import ast
import glob
import os
import shutil
import sys


def modify_toml(containerd_path):
    config_tomls = []
    
    with open(containerd_path, "r") as file:
        config_tomls = [line.rstrip() for line in file]

    for i, line in enumerate(config_tomls):
        if '[plugins."io.containerd.grpc.v1.cri".registry]' in line:
            config_tomls = config_tomls[:i+2]

    config_tomls[len(config_tomls) - 1] = '        config_path = "/etc/containerd/certs.d"'

    with open(containerd_path, "w") as file:
        for line in config_tomls:
            file.write(f"{line}\n")
        # file.write("\n".join(config_tomls))


def backup_toml(containerd_path):
    shutil.copyfile(containerd_path, containerd_path + ".bak")


def update_repos(containerd_path, repos: dict):
    dir = "certs.d"
    containerd_parent = os.path.dirname(containerd_path)
    for repo in repos:
        sub_dir = os.path.join(containerd_parent, dir, repo)
        if not os.path.exists(sub_dir):
            os.makedirs(sub_dir, exist_ok=True)
            print(f"Directory {sub_dir} created.")
        else:
            print(f"Directory {sub_dir} already exists.")
        
        sub_toml = os.path.join(sub_dir, "hosts.toml")
        with open(sub_toml, "w") as file:
            file.write(f"""[host."https://{repos.get(repo)}"]
  capabilities = ["pull", "resolve"]
""")

  
def check_repo(containerd_path, repos: dict):
    for repo in repos:
        sub_dir = os.path.join(containerd_path, dir, repo)
        if not os.path.exists(sub_dir):
            print(f"Repository {repo} does not exist.")
            exit(1)

        sub_toml = os.path.join(sub_dir, "hosts.toml")
        with open(sub_toml, "r") as file:
            for n, line in enumerate(file, start=1):
                line = line.strip()
                if n == 1:
                    if line != f'[host."https://{repos.get(repo)}"]':
                        print(f"Check repo fail in first line. {repo} in {dir}.")
                        print("{} : {}".format(line, f'[host."https://{repos.get(repo)}"]'))
                        exit(1)
                elif n == 2:
                    if line != f'  capabilities = ["pull", "resolve"]':
                        print(f"Check repo fail in second line. {repo} in {dir}.")


def display_repos(repo_path):
    files = [p for p in glob.glob(os.path.join(repo_path, "*", "*"), recursive=True) if os.path.isfile(p)]
    for f in files:
        with open(f, "r") as file:
            file_lines = file.readlines()
            print(f"File: {f}")
            for line in file_lines:
                print(line.strip())


def main():
    str_data = sys.argv[1] if sys.argv[1] else ""
    if str_data == "":
        print("Input error from repos.")
        exit(1)

    repos = ast.literal_eval(str_data)
    containerd_path = "/etc/containerd/config.toml"
    repo_path = "/etc/containerd/certs.d"
    backup_toml(containerd_path)
    modify_toml(containerd_path)
    update_repos(containerd_path, repos)
    display_repos(repo_path)
    # check_repo(containerd_path, repos)


if __name__ == "__main__":
    main()